from __future__ import absolute_import
import json

import logging
import traceback
from collections import namedtuple
from rest_framework import status
from rest_framework.response import Response
from rest_framework.views import APIView
from common.tools import record_ip, auth_request
from common.ploto_enum import CommonStatuEnum, TrainingJobEnum
from common.ploto_response import PlotoResponse
from common.train_tool import DataSet, Algorithm, Model, TrainingJob, Resource

logger = logging.getLogger(__name__)


class DataSetHandler(APIView):
    """
    数据集操作类
    """

    @record_ip
    @auth_request
    def get(self, request):
        """
        获取数据集的列表
        """
        response = PlotoResponse()
        data_set = DataSet()
        params = request.query_params.dict()
        page = int(params.get('page', 1))
        page_size = int(params.get('page_size', 10))
        name = params.get('dataset_name', "")
        remote_resp = data_set.get_datasets(page, page_size, name)
        text = json.loads(remote_resp.text)
        if remote_resp.status_code == 200:
            response.data['count'] = text['total_number']
            response.data['data_list'] = text['datasets']
            return Response(response.dict, status=status.HTTP_200_OK)
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        response.code = text['error_code']
        response.msg = text['error_msg']
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
        
    def post(self, request):
        """
        创建数据集
        """
        DataSetParam = \
        namedtuple('DataSetParam', ['data_path', 'dataset_name', 'description', 'work_path', 'dataset_type'])
        response = PlotoResponse()
        data_set = DataSet()
        if not request.user.has_perm('data_mgt.add_annotatedata'):
            response.code = CommonStatuEnum.AUTHORIZATION_EXCEPTION.code
            response.msg = CommonStatuEnum.AUTHORIZATION_EXCEPTION.msg
            return Response(response.dict, status=status.HTTP_403_FORBIDDEN)
        params = request.data
        try:
            param = DataSetParam(params['data_path'], params['dataset_name'], 
                                 params['description'], params['work_path'], params['dataset_type'])
        except Exception as e:
            logger.error("请求参数缺失:%s, %s, %s", params, repr(e), traceback.format_exc())
            response.code = CommonStatuEnum.VALID_EXCEPTION.code
            response.msg = CommonStatuEnum.VALID_EXCEPTION.msg
            return Response(response.dict, status=status.HTTP_400_BAD_REQUEST)
        remote_resp = data_set.add_dataset(param)
        if remote_resp.status_code == 201:
           return Response(response.dict, status=status.HTTP_201_CREATED)
        text = json.loads(remote_resp.text)
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        response.code = text['error_code']
        response.msg = text['error_msg']
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)

    def put(self, request):
        """
        单条数据集修改
        params: 数据集id、名、描述
        """
        params = request.data
        data_set = DataSet()
        response = PlotoResponse()
        if not request.user.has_perm('data_mgt.change_annotatedata'):
            response.code = CommonStatuEnum.AUTHORIZATION_EXCEPTION.code
            response.msg = CommonStatuEnum.AUTHORIZATION_EXCEPTION.msg
            return Response(response.dict, status=status.HTTP_403_FORBIDDEN)
        try:
            dataset_id = params['dataset_id']
            dataset_name = params['dataset_name']
            description = params['description']
        except Exception as e:
            logger.error("请求参数缺失:%s, %s, %s", params, repr(e), traceback.format_exc())
            response.code = CommonStatuEnum.VALID_EXCEPTION.code
            response.msg = CommonStatuEnum.VALID_EXCEPTION.msg
            return Response(response.dict, status=status.HTTP_400_BAD_REQUEST)
        remote_resp = data_set.update_dataset(dataset_id, dataset_name, description)
        if remote_resp.status_code == 200:
            return Response(response.dict, status=status.HTTP_200_OK)
        text = json.loads(remote_resp.text)
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        response.code = text['error_code']
        response.msg = text['error_msg']
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)

    def delete(self, request):
        # 删除数据集
        data_set = DataSet()
        response = PlotoResponse()
        if not request.user.has_perm('data_mgt.delete_annotatedata'):
            response.code = CommonStatuEnum.AUTHORIZATION_EXCEPTION.code
            response.msg = CommonStatuEnum.AUTHORIZATION_EXCEPTION.msg
            return Response(response.dict, status=status.HTTP_403_FORBIDDEN)
        try:
            dataset_id = request.data['dataset_id']
        except Exception as e:
            logger.error("数据集id错误:%s, %s, %s", request.data, repr(e), traceback.format_exc())
            response.code = CommonStatuEnum.VALID_EXCEPTION.code
            response.msg = CommonStatuEnum.VALID_EXCEPTION.msg
            return Response(response.dict, status=status.HTTP_400_BAD_REQUEST)
        remote_resp = data_set.delete_dataset(dataset_id=dataset_id)
        if remote_resp.status_code == 204:
            return Response(response.dict, status=status.HTTP_204_NO_CONTENT)
        text = json.loads(remote_resp.text)
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        response.code = text['error_code']
        response.msg = text['error_msg']
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)


class AnnotateHandler(APIView):
    """
    标注操作类
    """

    @record_ip
    @auth_request
    def post(self, request):
        """
        启动智能标注
        """
        params = request.data
        data_set = DataSet()
        response = PlotoResponse()
        annotate_type = params.get('annotate_type')
        dataset_id = params.get('dataset_id')
        if not annotate_type or not dataset_id \
                or (annotate_type != 'auto' and annotate_type != 'pre'):
            response.code = CommonStatuEnum.VALID_EXCEPTION.code
            response.msg = CommonStatuEnum.VALID_EXCEPTION.msg
            logger.error('请求参数错误%s', params)
            return Response(response.dict, status=status.HTTP_400_BAD_REQUEST)
        if annotate_type == 'auto':
            algorithm_type = params.get('algorithm_type', 'fast')
            remote_resp = data_set.auto_task(dataset_id, algorithm_type)
        else:
            model_id = params.get('model_id')
            specification = params.get('specification')
            instance_count = params.get('instance_count', 1)
            remote_resp = data_set.pre_task(dataset_id, model_id, specification, instance_count)
        if remote_resp.status_code == 200:
            return Response(response.dict, status=status.HTTP_200_OK)
        text = json.loads(remote_resp.text)
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        response.code = text['error_code']
        response.msg = text['error_msg']
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)


class AlgorithmHandeler(APIView):
    """
    算法操作类
    """

    @record_ip
    @auth_request
    def get(self, request):
        """
        获取算法列表（自建）
        """
        algorithm = Algorithm()
        response = PlotoResponse()
        params = request.query_params.dict()
        name = params.get('algorithm_name', "")
        page = int(params.get('page', 1))
        page_size = int(params.get('page_size', 10))
        remote_resp = algorithm.get_algorithms(page, page_size, name)
        text = json.loads(remote_resp.text)
        if remote_resp.status_code == 200:
            response.data['total'] = text['total']
            response.data['count'] = text['count']
            response.data['data_list'] = text['items']
            return Response(response.dict, status=status.HTTP_200_OK)
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        response.code = text['error_code']
        response.msg = text['error_msg']
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)

    def post(self, request):
        """
        新增算法
        支持通过预置框架和自定义两种方式创建
        """
        algorithm = Algorithm()
        response = PlotoResponse()
        params = request.data
        result = check_algorithm_param(params)
        if not result:
            response.code = CommonStatuEnum.VALID_EXCEPTION.code
            response.msg = CommonStatuEnum.VALID_EXCEPTION.msg
            logger.error('请求参数缺失%s', params)
            return Response(response.dict, status=status.HTTP_400_BAD_REQUEST)
        algorithms_detail = create_train_obj(params)
        remote_resp = algorithm.create_algorithms(algorithms_detail)
        text = json.loads(remote_resp.text)
        if remote_resp.status_code < 300:
            return Response(response.dict, status=status.HTTP_200_OK)
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        response.code = text['error_code']
        response.msg = text['error_msg']
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)

    def put(self, request):
        """
        更新算法名/描述
        """
        algorithm = Algorithm()
        response = PlotoResponse()
        params = request.data
        try:
            algorithm_id = params['algorithm_id']
            algorithm_name = params['algorithm_name']
            description = params['description']
        except Exception as e:
            logger.error("请求参数缺失:%s, %s, %s", params, repr(e), traceback.format_exc())
            response.code = CommonStatuEnum.VALID_EXCEPTION.code
            response.msg = CommonStatuEnum.VALID_EXCEPTION.msg
            return Response(response.dict, status=status.HTTP_400_BAD_REQUEST)
        remote_resp = algorithm.update_algorithms(algorithm_id, algorithm_name, description)
        if remote_resp.status_code < 300:
            return Response(response.dict, status=status.HTTP_200_OK)
        text = json.loads(remote_resp.text)
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        response.code = text['error_code']
        response.msg = text['error_msg']
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)

    def delete(self, request):
        """
        删除算法
        """
        algorithm = Algorithm()
        response = PlotoResponse()
        params = request.data
        algorithm_id = params.get('algorithm_id')
        if not algorithm_id:
            response.code = CommonStatuEnum.VALID_EXCEPTION.code
            response.msg = CommonStatuEnum.VALID_EXCEPTION.msg
            logger.error('算法id错误%s', params)
            return Response(response.dict, status=status.HTTP_400_BAD_REQUEST)
        remote_resp = algorithm.delete_algorithms(algorithm_id)
        if remote_resp.status_code < 300:
            return Response(response.dict, status=status.HTTP_204_NO_CONTENT)
        text = json.loads(remote_resp.text)
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        response.code = text['error_code']
        response.msg = text['error_msg']
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)


class AIHubHandeler(APIView):
    """
    AI Hub操作类
    """

    @record_ip
    @auth_request
    def get(self, request):
        """
        AI Gallery订阅信息
        """
        algorithm = Algorithm()
        model = Model()
        response = PlotoResponse()
        params = request.query_params.dict()
        type = params.get('type')
        if not type or (type != 'algo' and type != 'model'):
            logger.error('请求参数缺失%s', params)
            response.code = CommonStatuEnum.VALID_EXCEPTION.code
            response.msg = CommonStatuEnum.VALID_EXCEPTION.msg
            return Response(response.dict, status=status.HTTP_400_BAD_REQUEST)
        if type == 'algo':
            remote_resp = algorithm.get_hub_algorithms()
        else:
            remote_resp = model.get_hub_models()
        text = json.loads(remote_resp.text)
        if remote_resp.status_code == 200:
            response.data['count'] = text['total']
            response.data['data_list'] = text['subscription_infos']
            return Response(response.dict, status=status.HTTP_200_OK)
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        response.code = text['error_code']
        response.msg = text['error_msg']
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)


class TrainingJobHandler(APIView):
    """
    训练操作类
    """

    @record_ip
    @auth_request
    def get(self, request):
        """
        获取训练任务列表
        """
        training_job = TrainingJob()
        response = PlotoResponse()
        params = request.query_params.dict()
        name = params.get('job_name', "")
        page = int(params.get('page', 1))
        page_size = int(params.get('page_size', 10))
        remote_resp = training_job.get_training_job(page, page_size, name)
        text = json.loads(remote_resp.text)
        if remote_resp.status_code == 200:
            response.data['total'] = text['total']
            response.data['count'] = text['count']
            response.data['train_list'] = text['items']
            return Response(response.dict, status=status.HTTP_200_OK)
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        response.code = text['error_code']
        response.msg = text['error_msg']
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)

    def post(self, request):
        """
        创建训练任务
        """
        training_job = TrainingJob()
        response = PlotoResponse()
        params = request.data
        type = int(params.get('type', -1))
        if type not in (TrainingJobEnum.FRAME.code, TrainingJobEnum.CUSTOM.code,
                        TrainingJobEnum.ALGORITHM.code, TrainingJobEnum.AIHUB.code):
            response.code = CommonStatuEnum.VALID_EXCEPTION.code
            response.msg = CommonStatuEnum.VALID_EXCEPTION.msg
            logger.error('请求参数缺失%s', params)
            return Response(response.dict, status=status.HTTP_400_BAD_REQUEST)
        result = check_job_param(params)
        if not result:
            logger.error('请求参数缺失%s', params)
            response.code = CommonStatuEnum.VALID_EXCEPTION.code
            response.msg = CommonStatuEnum.VALID_EXCEPTION.msg
            return Response(response.dict, status=status.HTTP_400_BAD_REQUEST)
        job_detail = create_train_obj(params)
        remote_resp = training_job.create_training_jobs(job_detail, type)
        text = json.loads(remote_resp.text)
        if remote_resp.status_code < 300:
            return Response(response.dict, status=status.HTTP_200_OK)
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        response.code = text['error_code']
        response.msg = text['error_msg']
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)

    def put(self, request):
        """
        修改训练任务描述
        """
        training_job = TrainingJob()
        response = PlotoResponse()
        params = request.data
        try:
            job_id = params['job_id']
            description = params['description']
        except Exception as e:
            logger.error("请求参数缺失:%s, %s, %s", params, repr(e), traceback.format_exc())
            response.code = CommonStatuEnum.VALID_EXCEPTION.code
            response.msg = CommonStatuEnum.VALID_EXCEPTION.msg
            return Response(response.dict, status=status.HTTP_400_BAD_REQUEST)
        remote_resp = training_job.update_training_job(job_id, description)
        if remote_resp.status_code == 200:
            return Response(response.dict, status=status.HTTP_200_OK)
        text = json.loads(remote_resp.text)
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        response.code = text['error_code']
        response.msg = text['error_msg']
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)

    def delete(self, request):
        """
        删除训练任务
        """
        training_job = TrainingJob()
        response = PlotoResponse()
        params = request.data
        job_id = params.get('job_id')
        if not job_id:
            logger.error('训练id错误%s', params)
            response.code = CommonStatuEnum.VALID_EXCEPTION.code
            response.msg = CommonStatuEnum.VALID_EXCEPTION.msg
            return Response(response.dict, status=status.HTTP_400_BAD_REQUEST)
        remote_resp = training_job.delete_training_job(job_id)
        if remote_resp.status_code < 300:
            return Response(response.dict, status=status.HTTP_204_NO_CONTENT)
        text = json.loads(remote_resp.text)
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        response.code = text['error_code']
        response.msg = text['error_msg']
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)


class TrainingJobAction(APIView):
    """
    训练任务操作类
    """

    @record_ip
    @auth_request
    def put(self, request):
        """
        终止训练
        """
        training_job = TrainingJob()
        response = PlotoResponse()
        job_id = request.data.get('job_id')
        if not job_id:
            logger.error('训练id错误%s', job_id)
            response.code = CommonStatuEnum.VALID_EXCEPTION.code
            response.msg = CommonStatuEnum.VALID_EXCEPTION.msg
            return Response(response.dict, status=status.HTTP_400_BAD_REQUEST)
        remote_resp = training_job.stop_training_job(job_id)
        if remote_resp.status_code < 300:
            return Response(response.dict, status=status.HTTP_200_OK)
        text = json.loads(remote_resp.text)
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        response.code = text['error_code']
        response.msg = text['error_msg']
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)


class ModelHandler(APIView):
    """
    模型操作类
    """

    @record_ip
    @auth_request
    def get(self, request):
        """
        获取模型列表
        """
        model = Model()
        response = PlotoResponse()
        params = request.query_params.dict()
        page = int(params.get('page', 1))
        page_size = int(params.get('page_size', 10))
        name = params.get('model_name', "")
        remote_resp = model.get_models(page, page_size, name)
        text = json.loads(remote_resp.text)
        if remote_resp.status_code == 200:
            response.data['total'] = text['total_count']
            response.data['count'] = text['count']
            response.data['data_list'] = text['models']
            return Response(response.dict, status=status.HTTP_200_OK)
        response.code = text['error_code']
        response.msg = text['error_msg']
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)

    def post(self, request):
        """
        导入模型
        """
        model = Model()
        response = PlotoResponse()
        params = request.data
        if not check_model_param(params):
            logger.error('请求参数缺失%s', params)
            response.code = CommonStatuEnum.VALID_EXCEPTION.code
            response.msg = CommonStatuEnum.VALID_EXCEPTION.msg
            return Response(response.dict, status=status.HTTP_400_BAD_REQUEST)
        detail = create_train_obj(params)
        remote_resp = model.create_model(detail)
        if remote_resp.status_code < 300:
            return Response(response.dict, status=status.HTTP_200_OK)
        text = json.loads(remote_resp.text)
        response.code = text['error_code']
        response.msg = text['error_msg']
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)

    def delete(self, request):
        """
        删除模型
        """
        model = Model()
        response = PlotoResponse()
        model_id = request.data.get('model_id')
        if not model_id:
            logger.error('模型id错误%s', model_id)
            response.code = CommonStatuEnum.VALID_EXCEPTION.code
            response.msg = CommonStatuEnum.VALID_EXCEPTION.msg
            return Response(response.dict, status=status.HTTP_400_BAD_REQUEST)
        remote_resp = model.delete_model(model_id)
        text = json.loads(remote_resp.text)
        if remote_resp.status_code < 300:
            if not len(text.get('delete_failed_list')):
                return Response(response.dict, status=status.HTTP_204_NO_CONTENT)
            response.code = text['delete_failed_list'][0]['error_code']
            response.msg = text['delete_failed_list'][0]['error_msg']
            return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
        response.code = text['error_code']
        response.msg = text['error_msg']
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)


class FlavorHandler(APIView):
    """
    资源操作类
    """

    @record_ip
    @auth_request
    def get(self, request):
        """
        获取训练资源规格列表
        """
        resource = Resource()
        response = PlotoResponse()
        remote_resp = resource.get_flavors()
        text = json.loads(remote_resp.text)
        if remote_resp.status_code == 200:
            response.data['count'] = text['total_count']
            response.data['flavors_list'] = text['flavors']
            return Response(response.dict, status=status.HTTP_200_OK)
        response.code = text['error_code']
        response.msg = text['error_msg']
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)


class EngineHandler(APIView):
    """
    引擎操作类
    """

    @record_ip
    @auth_request
    def get(self, request):
        """
        获取AI引擎列表
        """
        resource = Resource()
        response = PlotoResponse()
        remote_resp = resource.get_engines()
        text = json.loads(remote_resp.text)
        if remote_resp.status_code == 200:
            response.data['count'] = text['total']
            response.data['engine_list'] = text['items']
            return Response(response.dict, status=status.HTTP_200_OK)
        logger.error("modelarts请求异常,%s,%s", text['error_code'], text['error_msg'])
        response.code = text['error_code']
        response.msg = text['error_msg']
        return Response(response.dict, status=status.HTTP_500_INTERNAL_SERVER_ERROR)


class TrainObj(object):
    def __init__(self, alog_dict):
        self.__dict__.update(alog_dict)


def create_train_obj(data):
    return json.loads(json.dumps(data), object_hook=TrainObj)


def check_algorithm_param(params):
    type = int(params.get('type'))
    name = params.get('name')
    if not all([name]) or type is None:
        return False
    # 预置框架
    if type == TrainingJobEnum.CUSTOM.code:
        image_url = params.get('image_url')
        command = params.get('command')
        if not all([image_url, command]):
            return False
        return True
    if type == TrainingJobEnum.FRAME.code:
        boot_file = params.get('boot_file')
        code_dir = params.get('code_dir')
        engine_name = params.get('engine_name')
        engine_version = params.get('engine_version')
        if not all([boot_file, code_dir, engine_name, engine_version]):
            return False
        return True
    return False


def check_job_param(params):
    type = int(params.get('type'))
    flavor_id = params.get('flavor_id')
    name = params.get('name')
    if not all([name, flavor_id]) or type is None:
        return False
    # 预置框架
    if type == TrainingJobEnum.FRAME.code:
        boot_file = params.get('boot_file')
        code_dir = params.get('code_dir')
        engine_name = params.get('engine_name')
        engine_version = params.get('engine_version')
        if not all([boot_file, code_dir, engine_name, engine_version]):
            return False
    if type == TrainingJobEnum.CUSTOM.code:
        image_url = params.get('image_url')
        command = params.get('command')
        if not all([image_url, command]):
            return False
    if type == TrainingJobEnum.ALGORITHM.code:
        id = params.get('id')
        if not id:
            return False
    if type == TrainingJobEnum.AIHUB.code:
        subscription_id = params.get('subscription_id')
        item_version_id = params.get('item_version_id')
        if not all([subscription_id, item_version_id]):
            return False
    return True


def check_model_param(params):
    model_name = params.get('model_name')
    model_version = params.get('model_version')
    source_location = params.get('source_location')
    model_type = params.get('model_type')
    runtime = params.get('runtime')
    if not all([model_name, model_type, model_version, source_location, runtime]):
        return False
    return True
